import numpy as np
import networkx as nx
import copy
SLOT_NUM=100
TASK_SLOT=10
TASK_REWARD=10
K=5

def create_nsfnet_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
    edges=[(0, 1), (0, 2), (0, 3), (1, 2), (1, 7), (2, 5), (3, 8), (3, 4), (4, 5), (4, 6), (5, 12), (5, 13),
         (6, 7), (7, 10), (8, 9), (8, 11), (9, 10), (9, 12), (10, 11), (10, 13), (11, 12)]
    Gbase.add_edges_from(edges)
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase
def create_geant2_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23])
    Gbase.add_edges_from(
        [(0, 1), (0, 2), (1, 3), (1, 6), (1, 9), (2, 3), (2, 4), (3, 6), (4, 7), (5, 3),
         (5, 8), (6, 9), (6, 8), (7, 11), (7, 8), (8, 11), (8, 20), (8, 17), (8, 18), (8, 12),
         (9, 10), (9, 13), (9, 12), (10, 13), (11, 20), (11, 14), (12, 13), (12,19), (12,21),
         (14, 15), (15, 16), (16, 17), (17,18), (18,21), (19, 23), (21,22), (22, 23)])
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase

def create_small_top():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8])
    Gbase.add_edges_from(
        [(0, 1), (0, 2), (0, 3), (1, 2), (1, 7), (2, 5), (3, 8), (3, 4), (4, 5), (4, 6), (5, 0),
         (6, 7), (6, 8), (7, 8), (8, 0), (8, 6), (3, 2), (5, 3)])
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase

def create_gbn_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
    Gbase.add_edges_from(
        [(0, 2), (0, 8), (1, 2), (1, 3), (1, 4), (2, 4), (3, 4), (3, 9), (4, 8), (4, 10), (4, 9),
         (5, 6), (5, 8), (6, 7), (7, 8), (7, 10), (9, 10), (9, 12), (10, 11), (10, 12), (11, 13),
         (12, 14), (12, 16), (13, 14), (14, 15), (15, 16)])
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase
def ksp(graph):
    path_record = dict()
    for n1 in graph:
        for n2 in graph:
            if n1 != n2:
                key = str(n1) + ':' + str(n2)
                if key not in path_record:
                    path_record[key] = []
                # cutoff 表示路由长度限制，超过此长度的路由不返回
                [path_record[key].append(p) for p in nx.all_simple_paths(graph, source=n1, target=n2, cutoff=14)]
                path_record[key] = sorted(path_record[key], key=lambda item: (len(item), item))
    for key,value in path_record.items():
        while len(value)<K:
            value.append(value[0])
    return path_record
# 计算边邻接关系，用作后续GAT的输入，隐式的进行了边节点转换
def cal_link_adj(graph):
    list1 = []
    list2 = []
    for u1, v1, d1 in (graph.edges(data=True)):
        for u2, v2, d2 in (graph.edges(data=True)):
            if (u1 == u2) | (u1 == v2) | (v1 == u2) | (v1 == v2):
                list1.append(int(d1['code']))
                list2.append(int(d2['code']))
    res = np.stack((list1, list2), axis=0)
    return res
def create_tasks(node_size,TASK_NUM):
    # 任务用一个三元组表示[head,tail,slot_num]
    res = np.zeros((TASK_NUM+2, 4), dtype=int)
    task_slots = np.random.poisson(TASK_SLOT, TASK_NUM+2)
    task_rewards=np.random.poisson(TASK_REWARD,TASK_NUM+2)
    for i in range(TASK_NUM+2):
        randoms = np.random.choice(np.array([p for p in range(node_size)]), 2, replace=False)
        head = randoms[0]
        tail = randoms[1]
        task_slot = task_slots[i]
        task_reward=task_rewards[i]
        res[i][0] = head
        res[i][1] = tail
        res[i][2] = task_slot
        res[i][3]=task_reward
    return res
# 输入路由节点list，输出路由链路编号
def get_link_code(graph, route):
    res = []
    for i in range(len(route) - 1):
        code = graph[int(route[i])][int(route[i + 1])]['code']
        res.append(code)
    return res
#查询连续0的个数和位置
#输入[0,1,1,0,0,0,1,0]
#输出[1,0,0,3,0,0,0,1]
def count_zero(list):
    res = np.zeros(len(list) + 1)
    index = 0
    num = 0
    for i in range(len(list)):
        if list[i] == 0:
            num += 1
        else:
            res[index] = num
            index = i + 1
            num = 0
    res[index] = num
    return res[0:len(list)]

#查询 路由 链路对应slot或运算的最终结果
# link是一个[链路数量，slot数量]的数组，route表示要查询的路由中链路编号
# 输出最终slot结果  [0,0,1,0,3,0,0,1,0]
def get_route_slot_final(link, route):
    link_res = np.zeros(len(link[0]))
    for i in route:
        for j in range(len(link[0])):
            link_res[j] = int(link_res[j]) | int(link[i][j])
    link_zero_res = count_zero(link_res)
    return link_zero_res
# 获取下一状态，以及是否分配成功  如果没有分配成功返回原状态和False
def get_next_state(curr_state1, link_code_list, task_slot):
    curr_state = copy.deepcopy(curr_state1)
    slot_available = get_route_slot_final(curr_state, link_code_list)
    # 当前路由缺少网络资源
    if np.max(slot_available) < task_slot:
        curr_state[0][0]=-1
        return curr_state, False
    else:
        index = np.inf
        for idx in range(len(slot_available)):
            if slot_available[idx] >= task_slot:
                index = idx
                break
        for i in link_code_list:
            for j in range(task_slot):
                curr_state[i][index + j] = 1
        return curr_state, True
class Network():
    def __init__(self,task_load):
        self.task_load=task_load
        self.graph=create_small_top()
        self.path_record=ksp(self.graph)
        self.link_adj=cal_link_adj(self.graph)
        self.node_size=nx.number_of_nodes(self.graph)
        self.link_size=nx.number_of_edges(self.graph)
        self.task=create_tasks(self.node_size,task_load)
        self.task_sum= np.sum(self.task[:,3])
        self.task_execute=0
        self.state_next=None

    def load_task(self):
        self.task=np.load('task.npy')
        self.task_sum = np.sum(self.task[:, 3])
        task_rewards=np.random.uniform(3,20,self.task_load+2)
        #task_rewards = np.random.poisson(TASK_REWARD, self.task_load + 2)
        # for i in range(self.task_load+2):
        #     self.task[i][3] = task_rewards[i]
        #     #self.task[i][3] = 1+i

    def reset(self):
        self.task_execute = 0
        for u, v, d in self.graph.edges(data=True):
            d['slots'] = np.zeros(SLOT_NUM)
    def observe(self):

        link_msg=np.zeros( (self.link_size,SLOT_NUM) )
        for u, v, d in self.graph.edges(data=True):
            link_msg[int(d['code'])] = d['slots']
        link_msg_next=np.zeros((K+1,self.link_size,SLOT_NUM))
        task=self.task[self.task_execute]
        for k in range(K):
            route = self.path_record[str(task[0]) + ':' + str(task[1])][k]
            route_code_list=get_link_code(self.graph,route)
            next_state,alloc_flag=get_next_state(link_msg,route_code_list,task[2])
            link_msg_next[k]=copy.deepcopy(next_state)
        link_msg_next[K]=copy.deepcopy(link_msg)
        self.state_next=link_msg_next
        return link_msg_next
    def step(self,action):
        if action==K:
            reward=0
        else:
            width = self.task[self.task_execute][3]/self.task_sum
            reward = width
            for u, v, d in self.graph.edges(data=True):
                d['slots'] = self.state_next[action][int(d['code'])][:SLOT_NUM]
        self.task_execute+=1
        return reward,self.task_execute>=self.task_load

create_tasks(14,100)
